import os
import json
from typing import Literal, TypedDict

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.language_models.chat_models import BaseChatModel # 导入 BaseChatModel
from langchain_openai import ChatOpenAI # 保留，将在 main.py 中使用
from langchain_community.chat_models import ChatTongyi # 保留，将在 main.py 中使用

# 如果使用通义千问，请取消注释以下行并注释掉 ChatOpenAI
# from langchain_community.chat_models import ChatTongyi

# 从环境变量加载 API Key
# 对于 OpenAI: OPENAI_API_KEY
# 对于通义千问: DASHSCOPE_API_KEY
# 请确保在 .env 文件中设置了相应的环境变量
# 例如：OPENAI_API_KEY="sk-..." 或 DASHSCOPE_API_KEY="sk-..."

# 全局 LLM 实例已移除，将在 main.py 中根据配置动态创建

# 定义输出类型，用于 Langchain 的 with_structured_output
class QuestionItem(TypedDict):
    id: Literal["character", "task", "context", "format", "constraints", "tone", "examples"]
    question: str

class QuestionsOutput(TypedDict):
    type: Literal["questions"]
    content: list[QuestionItem]

class OptimizedPromptOutput(TypedDict):
    type: Literal["optimized_prompt"]
    content: str

def _read_prompt_file(file_name: str) -> str:
    """Helper function to read prompt content from a Markdown file."""
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # Go up to 'backend' directory
    prompts_dir = os.path.join(base_dir, "promotes") # Navigate to 'promotes' directory
    file_path = os.path.join(prompts_dir, file_name)
    with open(file_path, 'r', encoding='utf-8') as f:
        return f.read()

# 1. 分析与提问模板 (Analyzer/Questioner Prompt)
# 这个模板用于分析用户输入，并根据 CRF-CT-TE 框架提出一系列问题。
analyzer_questioner_system_prompt = _read_prompt_file("analyzer_questioner_system_prompt.md")

analyzer_questioner_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", analyzer_questioner_system_prompt),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}"),
    ]
)

# 2. 最终优化模板 (Optimizer Prompt)
# 这个模板用于根据完整的对话历史（原始提示词 + 所有问答）生成最终的优化提示词。
optimizer_system_prompt = _read_prompt_file("optimizer_system_prompt.md")

optimizer_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", optimizer_system_prompt),
        MessagesPlaceholder(variable_name="history"),
    ]
)

# 创建 Langchain 链的函数
# 这些函数接收一个 LLM 实例作为参数

def get_analyzer_questioner_chain(llm: BaseChatModel):
    """
    构建并返回分析与提问链。
    :param llm: 一个 BaseChatModel 实例 (例如 ChatOpenAI, ChatTongyi)。
    :return: 配置好的 Langchain 链。
    """
    return (
        analyzer_questioner_prompt
        | llm
        | JsonOutputParser()
    )

def get_optimizer_chain(llm: BaseChatModel):
    """
    构建并返回最终优化链。
    :param llm: 一个 BaseChatModel 实例 (例如 ChatOpenAI, ChatTongyi)。
    :return: 配置好的 Langchain 链。
    """
    return (
        optimizer_prompt
        | llm
        | JsonOutputParser()
    )
